#include <stdio.h>
#include <stdlib.h>
#include <openssl/bn.h>
#include <string.h>
#include "code_table.h"

typedef struct key{
    BIGNUM *exp;
    BIGNUM *N;
}Key;

// 大素数选择返回参数（公钥指数，私钥指数，N）的结构体
typedef struct rsa
{
    BIGNUM *e;
    BIGNUM *d;
    BIGNUM *N;
}RSA_t;

// 设置大数值
BIGNUM *Set(char *num);

// 打印大数
void PrintBigInt(BIGNUM *a);
BIGNUM *BigIntadd(BIGNUM *a, char *b);
// 大数减法封装
BIGNUM *BigIntsub(BIGNUM *a, char *b);
// 计算欧拉
BIGNUM *Eular(BIGNUM *a, BIGNUM *b);
// 计算gcd

RSA_t SelectPrime(int kbits, int ebits);

Key gennerate_public_key(RSA_t r);
Key gennerate_private_key(RSA_t r);

BIGNUM **Encryption_t(char **Plaintext, int clen, Key public_key);
BIGNUM **Decryption_t(BIGNUM **ciphertext, int clen, Key private_key);

int main()
{   
    int i;

    // 确定k和公钥指数e的位数
    int kbits;
    int ebits;
    printf("input k and exp for public_key\n");
    scanf("%d %d", &kbits, &ebits);

    // 大素数选择流程
    RSA_t rsa;
    rsa = SelectPrime(kbits, ebits);

    Key pubKey = gennerate_public_key(rsa);
    Key priKey = gennerate_private_key(rsa);

    printf("\npublicKey.e privateKey.d Key.N\n");
    PrintBigInt(rsa.e);
    printf("\n");
    PrintBigInt(rsa.d);
    printf("\n");
    PrintBigInt(rsa.N);
    printf("\n");
    
    // 编解码表生成
    Code_table ct = initCode_Table();
    
    // 输入明文
    printf("\ninput PlainText\n");
    char *platex;
    platex = (char *)malloc(10 * sizeof(char));
    scanf("%s", platex);
    
    // 明文字符串化
    char *platext[10] = {"0"};
    int clen = strlen(platex);
    for(i = 0; i < clen; i++){
        char *temp;
        temp = (char *)malloc(10 * sizeof(char));
        //printf("%c\n", platex[i]);  
        sprintf(temp, "%c", platex[i]);
        platext[i] = temp;
    }

    // 编解码转换
    char *code[10] = {"0"};
    for(i = 0; i < clen; i++){
        code[i] = lettertoCode(platext[i], ct);
    }
    
    // 加密
    printf("\nEncrypt...\n\n");
    BIGNUM **ciptext = Encryption_t(code, clen, pubKey);
    printf("CipherText string\n");
    for(i = 0; i < clen; i++){
        int temp = i;
        printf("C%d: ", temp+1);
        PrintBigInt(ciptext[i]);
        printf("\n");
    }
    
    // 解密
    char *plaintex;
    printf("\nDecrypt...\n\n");
    BIGNUM **text = Decryption_t(ciptext, clen, priKey);
    printf("PlainText\n");
    for(i = 0; i < clen; i++){
        int temp = i;
        plaintex = BN_bn2dec(text[i]);
        plaintex = codeToletter(plaintex, ct);
        printf("M%d: %s \n", temp+1, plaintex);
    }
    printf("\n");
    return 0;
}

Code_table initCode_Table(){
    Code_table ct;
    int i;
    for(i = 0; i < 26; i++){
        int temp = i;
        char *ctr;
        ctr = (char *)malloc(10 * sizeof(char));
        temp++;
        ct.letter[i] = letter[i]; 
        sprintf(ctr, "%d", temp);
        ct.code[i] = ctr;
    }
    return ct;
}

char *codeToletter(char *code, Code_table ct){
    int i;
    for(i = 0; i < 26; i++){
        if(strcmp(code, ct.code[i]) == 0){
            return ct.letter[i];
        }
    }
}

char *lettertoCode(char *letter, Code_table ct){
    int i;
    for(i = 0; i < 26; i++){
        if(strcmp(letter, ct.letter[i]) == 0){
            return ct.code[i];
        }
    }
}

BIGNUM *Set(char *num){
    BIGNUM *res;
    res = BN_new();
    BN_dec2bn(&res, num);
    return res;
    free(res);
}

void PrintBigInt(BIGNUM *a){
    char *c;
    c = BN_bn2hex(a);
    //puts(c);
    printf("%s", c);
    free(c);
}

BIGNUM *BigIntadd(BIGNUM *a, char *b){
    BIGNUM *bn, *res;     // bn：大数类型的b, res: 和
    bn = Set(b);        //将字符十进制被减数转换为大数类型
    res = BN_new();
    BN_add(res, a, bn);
    return res;
}

BIGNUM *BigIntsub(BIGNUM *a, char *b){
    BIGNUM *bn, *res;     // bn：大数类型的b, res: 和
    bn = Set(b);        //将字符十进制被减数转换为大数类型
    res = BN_new();
    BN_sub(res, a, bn);
    return res;
}

BIGNUM *Eular(BIGNUM *a, BIGNUM *b){
    BIGNUM *res, *a_1, *b_1;
    res = BN_new();
    a_1 = BN_new();
    b_1 = BN_new();
    a_1 = BigIntsub(a, "1");
    b_1 = BigIntsub(b, "1");
    BN_CTX *ctx;
    ctx = BN_CTX_new();
    BN_mul(res, a_1, b_1, ctx);
    return res;
}

RSA_t SelectPrime(int kbits, int ebits)
{
    RSA_t rsa;
    
    BIGNUM *p, *q, *N, *e, *d, *eul;
    //step1 确定位数
    int k = kbits;
    int pbits = (k + 1) / 2;
    int qbits = (k - 1) / 2;
    p = BN_new();
    q = BN_new();
    N = BN_new();
    e = BN_new();
    d = BN_new();
    eul = BN_new();

    //step2 生成素数
    BN_generate_prime_ex(p, pbits, 10, NULL, NULL, NULL);
    BN_generate_prime_ex(q, qbits, 10, NULL, NULL, NULL);
    

    //step3 计算N=pq
    BN_CTX *ctx;
    ctx = BN_CTX_new();
    BN_mul(N, p, q, ctx);
    
    //step4 选择公钥指数e
    eul = Eular(p, q);
    while (1)
    {
        BN_rand(e, ebits, 0, 0);
        BIGNUM *gcd, *Bigone;
        gcd = BN_new();
        Bigone = BN_new();
        char one = '1';
        BN_dec2bn(&Bigone, &one);
        BN_gcd(gcd, e, eul, ctx);
        if (BN_cmp(gcd, Bigone) == 0)
        {
            break;
        }
    }
    
    //step5 求秘钥指数d d = (eul * i + 1) / e ... mod = 0
    BIGNUM *i = Set("1");
    while (1)
    {
        BIGNUM *ieul, *mod;
        ieul = BN_new();
        mod = BN_new();
        BN_mul(ieul, eul, i, ctx);
        ieul = BigIntadd(ieul, "1");
        BN_div(d, mod, ieul, e, ctx);
        
        if(BN_cmp(mod, Set("0")) == 0){
            break;
        }
        i = BigIntadd(i, "1");
    }
    
    rsa.e = e;
    rsa.d = d;
    rsa.N = N;
    return rsa;
}

Key gennerate_public_key(RSA_t r){
    Key public_k;
    public_k.exp = r.e;
    public_k.N = r.N;
    return public_k;
}

Key gennerate_private_key(RSA_t r){
    Key private_k;
    private_k.exp = r.d;
    private_k.N = r.N;
    return private_k;
}

BIGNUM **Encryption_t(char **Plaintext, int clen, Key public_key){
    BIGNUM **ciptext;
    ciptext = BN_new();
    int i;

    for(i = 0; i < clen; i++){
        BIGNUM *temp;
        temp = BN_new();
        temp = Set(Plaintext[i]);
        BN_CTX *ctx;
        ctx = BN_CTX_new();
        
        // Ci = Mi ^ e mod N
        BN_mod_exp(temp, temp, public_key.exp, public_key.N, ctx);
        ciptext[i] = temp;        
        BN_copy(ciptext[i], temp);        
    }

    return ciptext;
}

BIGNUM **Decryption_t(BIGNUM **ciphertext, int clen, Key private_key){
    BIGNUM **platext;
    platext = BN_new();
    int i;
    
    for(i = 0; i < clen; i++){
        BIGNUM *temp;
        temp = BN_new();
        temp = ciphertext[i];
        BN_CTX *ctx;
        ctx = BN_CTX_new();
        // Mi = Ci ^ d mod N
        BN_mod_exp(temp, temp, private_key.exp, private_key.N, ctx);
        platext[i] = temp;
    }
    return platext;
}
